#!/usr/bin/env python3
"""
Kafka测试数据生成脚本
生成10万条测试消息到test_kafka主题
"""

import json
import random
import time
from datetime import datetime, timedelta
from typing import List, Dict, Any
import asyncio
from kafka import KafkaProducer
from kafka.errors import KafkaError
import logging

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Kafka配置
KAFKA_CONFIG = {
    'bootstrap_servers': ['192.168.37.11:9092'],
    'value_serializer': lambda v: json.dumps(v, ensure_ascii=False).encode('utf-8'),
    'key_serializer': lambda k: str(k).encode('utf-8') if k else None,
    'acks': 1,  # 等待leader确认
    'retries': 3,
    'batch_size': 16384,
    'linger_ms': 10,  # 批量发送延迟
    'buffer_memory': 33554432,
    'compression_type': 'gzip'
}

# 测试数据模板
MESSAGE_TEMPLATES = [
    {
        "type": "message",
        "content": "密码将在 {} 天后过期",
        "status": "已送达",
        "category": "security"
    },
    {
        "type": "message", 
        "content": "API 调用频率超限",
        "status": "已读",
        "category": "api"
    },
    {
        "type": "message",
        "content": "备份任务执行成功",
        "status": "已发送",
        "category": "system"
    },
    {
        "type": "notification",
        "content": "新用户注册: {}",
        "status": "已发送",
        "category": "user"
    },
    {
        "type": "alert",
        "content": "服务器CPU使用率达到 {}%",
        "status": "已送达",
        "category": "monitoring"
    },
    {
        "type": "message",
        "content": "订单 {} 已完成支付",
        "status": "已读",
        "category": "order"
    },
    {
        "type": "notification",
        "content": "系统维护将在 {} 小时后开始",
        "status": "已发送",
        "category": "maintenance"
    },
    {
        "type": "alert",
        "content": "数据库连接异常",
        "status": "已送达",
        "category": "database"
    },
    {
        "type": "message",
        "content": "文件上传完成: {}",
        "status": "已读",
        "category": "file"
    },
    {
        "type": "notification",
        "content": "邮件发送成功: {} 封",
        "status": "已发送",
        "category": "email"
    }
]

# 状态选项
STATUSES = ["已发送", "已送达", "已读", "失败", "处理中"]

# 类型选项
MESSAGE_TYPES = ["message", "notification", "alert", "warning", "info"]

# 内容变量
CONTENT_VARIABLES = {
    "days": [1, 2, 3, 5, 7, 10, 15, 30],
    "usernames": ["user001", "admin", "guest", "test_user", "developer", "manager"],
    "percentages": [75, 80, 85, 90, 95, 98],
    "order_ids": ["ORD001", "ORD002", "ORD003", "PAY001", "PAY002"],
    "hours": [1, 2, 4, 6, 8, 12, 24],
    "filenames": ["data.csv", "backup.zip", "report.pdf", "image.jpg", "config.json"],
    "counts": [10, 25, 50, 100, 200, 500, 1000]
}

class KafkaTestDataGenerator:
    def __init__(self, topic_name: str = "test_kafka", total_messages: int = 100000):
        self.topic_name = topic_name
        self.total_messages = total_messages
        self.producer = None
        self.sent_count = 0
        self.failed_count = 0
        
    def create_producer(self) -> KafkaProducer:
        """创建Kafka生产者"""
        try:
            producer = KafkaProducer(**KAFKA_CONFIG)
            logger.info("Kafka生产者创建成功")
            return producer
        except Exception as e:
            logger.error(f"创建Kafka生产者失败: {e}")
            raise
    
    def generate_message(self, message_id: int) -> Dict[str, Any]:
        """生成单条测息"""
        # 随机选择模板
        template = random.choice(MESSAGE_TEMPLATES)
        
        # 生成基础消息
        message = {
            "id": message_id,
            "type": template["type"],
            "content": template["content"],
            "status": random.choice(STATUSES),
            "category": template.get("category", "general")
        }
        
        # 处理内容中的占位符
        if "{}" in message["content"]:
            if "天" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["days"]))
            elif "%" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["percentages"]))
            elif "用户" in message["content"] or "注册" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["usernames"]))
            elif "订单" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["order_ids"]))
            elif "小时" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["hours"]))
            elif "文件" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["filenames"]))
            elif "封" in message["content"]:
                message["content"] = message["content"].format(random.choice(CONTENT_VARIABLES["counts"]))
        
        # 生成时间戳 (最近7天内的随机时间)
        base_time = datetime.now()
        random_days = random.randint(0, 7)
        random_hours = random.randint(0, 23)
        random_minutes = random.randint(0, 59)
        random_seconds = random.randint(0, 59)
        
        message_time = base_time - timedelta(
            days=random_days,
            hours=random_hours,
            minutes=random_minutes,
            seconds=random_seconds
        )
        
        message["timestamp"] = message_time.strftime("%Y-%m-%dT%H:%M:%SZ")
        
        # 添加一些随机字段
        if random.random() < 0.3:  # 30%概率添加额外字段
            message["priority"] = random.choice(["low", "medium", "high", "urgent"])
        
        if random.random() < 0.2:  # 20%概率添加用户ID
            message["user_id"] = f"user_{random.randint(1000, 9999)}"
        
        if random.random() < 0.1:  # 10%概率添加错误码
            message["error_code"] = random.choice([200, 400, 401, 403, 404, 500, 502, 503])
        
        return message
    
    def send_message_callback(self, message_id: int):
        """发送消息的回调函数"""
        def callback(record_metadata=None, exception=None):
            if exception:
                self.failed_count += 1
                logger.error(f"消息 {message_id} 发送失败: {exception}")
            else:
                self.sent_count += 1
                if self.sent_count % 1000 == 0:
                    logger.info(f"已发送 {self.sent_count}/{self.total_messages} 条消息")
        return callback
    
    def generate_and_send_batch(self, start_id: int, batch_size: int):
        """生成并发送一批消息"""
        for i in range(batch_size):
            message_id = start_id + i
            if message_id > self.total_messages:
                break
                
            # 生成消息
            message = self.generate_message(message_id)
            
            # 生成Key (使用用户ID或消息ID)
            key = message.get("user_id", f"msg_{message_id}")
            
            try:
                # 异步发送消息
                future = self.producer.send(
                    self.topic_name,
                    value=message,
                    key=key
                )
                future.add_callback(self.send_message_callback(message_id))
                
            except Exception as e:
                logger.error(f"发送消息 {message_id} 时出错: {e}")
                self.failed_count += 1
    
    def run(self):
        """运行数据生成"""
        logger.info(f"开始生成 {self.total_messages} 条测试消息到主题: {self.topic_name}")
        
        try:
            # 创建生产者
            self.producer = self.create_producer()
            
            # 批量发送配置
            batch_size = 1000
            start_time = time.time()
            
            # 分批发送消息
            for batch_start in range(1, self.total_messages + 1, batch_size):
                current_batch_size = min(batch_size, self.total_messages - batch_start + 1)
                self.generate_and_send_batch(batch_start, current_batch_size)
                
                # 每10批次刷新一次
                if (batch_start // batch_size) % 10 == 0:
                    self.producer.flush()
                    time.sleep(0.1)  # 短暂休息，避免过载
            
            # 确保所有消息都发送完成
            logger.info("等待所有消息发送完成...")
            self.producer.flush()
            
            # 统计结果
            end_time = time.time()
            duration = end_time - start_time
            
            logger.info("=" * 50)
            logger.info("数据生成完成!")
            logger.info(f"总消息数: {self.total_messages}")
            logger.info(f"成功发送: {self.sent_count}")
            logger.info(f"发送失败: {self.failed_count}")
            logger.info(f"耗时: {duration:.2f} 秒")
            logger.info(f"平均速度: {self.sent_count/duration:.2f} 消息/秒")
            logger.info("=" * 50)
            
        except KeyboardInterrupt:
            logger.info("用户中断，正在停止...")
        except Exception as e:
            logger.error(f"生成数据时出错: {e}")
        finally:
            if self.producer:
                self.producer.close()
                logger.info("Kafka生产者已关闭")

def main():
    """主函数"""
    import argparse
    
    parser = argparse.ArgumentParser(description='Kafka测试数据生成器')
    parser.add_argument('--topic', default='test_kafka', help='目标主题名称')
    parser.add_argument('--count', type=int, default=100000, help='生成消息数量')
    parser.add_argument('--server', default='192.168.37.11:9092', help='Kafka服务器地址')
    
    args = parser.parse_args()
    
    # 更新Kafka配置
    if args.server != '192.168.37.11:9092':
        KAFKA_CONFIG['bootstrap_servers'] = [args.server]
    
    # 创建生成器并运行
    generator = KafkaTestDataGenerator(
        topic_name=args.topic,
        total_messages=args.count
    )
    
    generator.run()

if __name__ == "__main__":
    main()
